#include "core/memory.h"
#include "comm/types.h"
#include "ipc/mutex.h"
#include "tools/log.h"
#include "tools/klib.h"
#include "cpu/mmu.h"
#include "dev/console.h"

// 注意静态变量位于BSS段自动清零
static addr_alloc_t paddr_alloc;        // 内存也空间分配器
static pde_t kernel_page_dir[PDE_CNT] __attribute__((aligned(MEM_PAGE_SIZE)));  // 内核页目录表

// 内存空间分配器初始化
static void addr_alloc_init(addr_alloc_t *alloc, uint8_t *bits,
    uint32_t start, uint32_t size, uint32_t page_size) {
    mutex_init(&alloc->mutex);
    alloc->start = start;
    alloc->size = size;
    alloc->page_size = page_size;
    bitmap_init(&alloc->bitmap, bits, alloc->size / page_size, 0);
}

// 分配连续的page_count页内存(申请一页内存可以用作页表，当然也可以用作别的)
// 返回值为申请的内存页的首地址
static uint32_t addr_alloc_page(addr_alloc_t *alloc, int page_count) {
    uint32_t addr = 0;
    mutex_lock(&alloc->mutex);
    int page_index = bitmap_alloc_nbits(&alloc->bitmap, 0, page_count);
    if (page_index >= 0) {
        addr = alloc->start + page_index * alloc->page_size;
    }
    mutex_unlock(&alloc->mutex);
    return addr;
}

// 释放连续的page_count个内存页
static void addr_free_page(addr_alloc_t *alloc, uint32_t addr, int page_count) {
    mutex_lock(&alloc->mutex);
    uint32_t page_index = (addr - alloc->start) / alloc->page_size;
    bitmap_set_bit(&alloc->bitmap, page_index, page_count, 0);
    mutex_unlock(&alloc->mutex);
}

// 展示从BIOS中获取的可以使用的内存段信息
void show_mem_info(boot_info_t *boot_info) {
    log_printf("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i ++) {
        log_printf("[%d]: 0x%x-0x%x", i,
        boot_info->ram_region_cfg[i].start,
        boot_info->ram_region_cfg[i].size );
    }
}

// 计算从BIOS中获取的可以使用的内存段信息总字节数
static uint32_t total_mem_size(boot_info_t *boot_info) {
    uint32_t mem_size = 0;
    for (int i = 0; i < boot_info->ram_region_count; i ++) {
        mem_size += boot_info->ram_region_cfg[i].size;
    }
    return mem_size;
}


// 概念解析：
// pde表示的是页目录表page_dir中的页目录表项
// pte表示的是页表page_table中的页表项
// page_dir中的每个表项pde都指向一个page_table
// 每个page_table中又含有1024个pte
// 内核使用的页目录表是kernel_page_dir

// 从页目录表中，根据vaddr的高10位去找对应的pde，
// 在根据ped去找它指向的页表page_table，
// 如果页表不存在，那么就根据alloc字段的值，去决定是否去分配一页内存作为page_table，
// 如果alloc为1我们就分配一页内存作为page_table，并且设置pde表项的字段值令其指向新分配的page_table
// 在找到了/新分配了page_table后，根据vaddr的中间10位去找位于page_table中的页表项pte，
// 找到pte后我们返回这个pte的收地址(以便让别的函数去修改这个pte表项的的各个字段值(指向的实际物理地址、访问这个物理地址的权限等))
pte_t * find_pte(pde_t *page_dir, uint32_t vaddr, int alloc) {
    pte_t * page_table;
    pde_t *pde = page_dir + pde_index(vaddr);

    // 如果page_dir中的页目录表项pde存在，说明这个表项所指向的页表是存在的，那么我们就直接赋值
    if (pde->present) {
        page_table = (pte_t *)pde_addr(pde);
    } else {
        // pde指向的页表不存在且不让分配内页表为这个ped
        if (alloc == 0) {
            return (pte_t *)0;
        }
        // pde指向的页表不存，我们新分配一页内存作为页表pg_table，并让当前这个页目录表项指向pg_table起始地址pa_paddr
        uint32_t pg_paddr = addr_alloc_page(&paddr_alloc, 1);
        if (pg_paddr == 0) {
            return (pte_t *)0;
        }
        // 给当前这个页目录表项设置指向的页表的权限位和起始地址
        pde->v = pg_paddr | PDE_P | PDE_W | PDE_U;

        page_table = (pte_t *)pg_paddr;
        // 清空分配给我的页表
        kernel_memset(page_table, 0, MEM_PAGE_SIZE);
    }
    return page_table + pte_index(vaddr);
}

// 将指定的一段地址空间建立虚拟地址vaddr到物理地址paddr的映射映射关系
// 映射关系是存储在参数page_dir指定的页目录表中
// 这段空间占用count个内存页，访问这段空间的属性是perm
int memory_create_map(pde_t * page_dir, uint32_t vaddr, uint32_t paddr, int count, int perm) {
    for (int i = 0; i < count; i++) {
        // log_printf("create map: v-0x%x, p-0x%x, perm:0x%x", vaddr, paddr, perm);
        
        pte_t *pte = find_pte(page_dir, vaddr, 1);      // 1表示没有找到相应的pte表就分配一个pte表
        if (pte == (pte_t *)0) {
            log_printf("create pte failed. pte == 0");
            return -1;
        }

        // log_printf("\tpte addr:0x%x", (uint32_t)pte);
        ASSERT(pte->present == 0);      // 目前我们在建立映射关系，说明之前是不存在这个映射关系的
        
        pte->v = paddr | perm | PTE_P;
        
        vaddr += MEM_PAGE_SIZE;
        paddr += MEM_PAGE_SIZE;
    }
    return 0;
}

// 将内核程序的所在的空间分成几段，对每一段分别建立虚拟地址空间到物理地址空间到映射
void create_kernel_table(void) {
    extern uint8_t kernel_base[];
    extern uint8_t s_text[], e_text[], s_data[];
    
    static memory_map_t kernel_map[] = {
        { kernel_base, s_text, kernel_base, PTE_W },        // 内核栈区
        { s_text, e_text, s_text, 0 },                      // 内核代码区，只读不用增加权限
        {(void *)CONSOLE_DISP_ADDR, (void *)(CONSOLE_DISP_END - 1), (void *)CONSOLE_VIDEO_BASE, PTE_W}, // 显存段映射
        { s_data, (void *)MEM_EBDA_START, s_data, PTE_W },  // 内核数据区
        
        // 把1MB以上的内存区也建立一个映射，便于内核以"虚拟地址==物理地址"的方式进行直接操作
        { (void *)MEM_EXT_START, (void *)MEM_EXT_END, (void *)MEM_EXT_START, PTE_W },
    }; 

    // 清空页目录表(可省略，静态变量默认自动清零)
    kernel_memset(kernel_page_dir, 0, sizeof(kernel_page_dir));

    // 根据映射关系依次创建映射表(kernel_page_dir)
    for (int i = 0; i < sizeof(kernel_map) / sizeof(memory_map_t); i++) {
        memory_map_t *map = kernel_map + i;

        // 可能有多个页，建立多个页的配置
        // 简化起见，不考虑4M的情况
        uint32_t vstart = down2((uint32_t)map->vstart, MEM_PAGE_SIZE);
        uint32_t pstart = down2((uint32_t)map->pstart, MEM_PAGE_SIZE);
        uint32_t vend = up2((uint32_t)map->vend, MEM_PAGE_SIZE);
        int page_count = (vend - vstart) / MEM_PAGE_SIZE;

        // 将一定地址空间进行映射到kernel_page_dir中
        memory_create_map(kernel_page_dir, vstart, pstart, page_count, map->perm); 
    }
}

// 初始化内存管理系统和创建内核地址空间映射
// 1. 初始化物理内存分配器：将所有的物理内存管理起来。在1MB内存中分配内存页物理位图
// 2. 虽然之前已经通过enable_page_mode函数建立了一个单级分页，并映射了虚拟存在从0开始到4mb到物理内存的0开始处的4mb
// 对于映射的这个段权限不够细分，因此这里我们同事使用二级页表的方式，重新对内核程序的多个地址空间段进行映射，对每个地址空间段作更精细的访问权限控制
void memory_init(boot_info_t *boot_info) {
    // 1MB内存空间起始，在链接脚本中定义
    extern uint32_t mem_free_start;         // 该变量本质是lds指定在汇编中的标号，标号的地址就是本变量的地址
    
    show_mem_info(boot_info);

    // 在内核数据后面放物理页位图
    uint8_t *mem_free = (uint8_t *)&mem_free_start;     // 用来存放内存页位图

    // 计算1MB以上空间的空闲内存容量，并对齐的页边界
    uint32_t mem_up1MB_free = total_mem_size(boot_info) - MEM_EXT_START;
    mem_up1MB_free = down2(mem_up1MB_free, MEM_PAGE_SIZE);
    // log_printf("free memory: 0x%x, 0x%x", MEM_EXT_START, mem_up1MB_free);
 
    // 4GB大小需要总共4*1024*1024*1024/4096/8=128KB的位图, 使用低1MB的RAM空间中足够
    // 该部分的内存仅跟在mem_free_start开始放置
    addr_alloc_init(&paddr_alloc, mem_free, MEM_EXT_START, mem_up1MB_free, MEM_PAGE_SIZE);
    mem_free += bitmap_byte_count(paddr_alloc.size / MEM_PAGE_SIZE);
    // 到这里mem_free应该比EBDA的地址要小
    ASSERT(mem_free <= (uint8_t *)MEM_EBDA_START);
    // 创建内核地址空间映射表
    create_kernel_table();
    // 设置内核页目录映射表地址到CR3中
    mmu_set_page_dir((uint32_t)kernel_page_dir);
}


// 创建进程的初始页表
// 主要的工作创建页目录表，然后从内核页表中复制一部分
// 注意复制内核页表kernel_page_dir的原因：
// 内核程序是所有应用进程共享的，当用户程序使用某个系统调用的时候就需要用到内核页表的映射关系，
// 通过使用这个映射关系来完成系统调用函数的虚拟地址到物理地址的转换
uint32_t memory_create_uvm(void) {
    pde_t *page_dir = (pde_t *)addr_alloc_page(&paddr_alloc, 1);
    if (page_dir == 0) {
        return 0;
    }
    kernel_memset((void *)page_dir, 0, MEM_PAGE_SIZE);
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    for (int i = 0; i < user_pde_start; i ++) {
        page_dir[i].v = kernel_page_dir[i].v;
    }

    return (uint32_t)page_dir;
}

// 为某个进程的page_dir的虚拟地址空间分配物理内存并建立映射关系
int memory_alloc_for_page_dir(uint32_t page_dir, uint32_t vaddr, uint32_t size, int perm) {
    uint32_t curr_vaddr = vaddr;
    int page_count = up2(size, MEM_PAGE_SIZE) / MEM_PAGE_SIZE;

    for (int i = 0; i < page_count; i ++) {
        // 首先分配物理内存
        uint32_t paddr = addr_alloc_page(&paddr_alloc, 1);
        if (paddr == 0) {
            log_printf("mem alloc failed. no memory");
            return 0;
        }
        // 建立虚拟内存和物理内存的映射关系
        int err = memory_create_map((pde_t *)page_dir, curr_vaddr, paddr, 1, perm);
        if (err < 0) {
            log_printf("create memory failed. err = %d", err);
            addr_free_page(&paddr_alloc, vaddr, page_count);
            // 这里需要考虑清除page_dir和page_table中的映射关系
            // 此处函数还没有实现
            return 0;
        }
        curr_vaddr += MEM_PAGE_SIZE;
    }

    return 0;
}

// 为进程分配页面内存
int memory_alloc_page_for(uint32_t addr, uint32_t size, uint32_t perm) {
    return memory_alloc_for_page_dir(task_current()->tss.cr3, addr, size, perm);
}

// 为内核分配1页内存
uint32_t memory_alloc_page() {
    // 注意addr_alloc_page函数分配的是物理地址
    // 由于我们对kernel_map 1mb以上的内存区域做了映射
    // 也就是说在1mb以上的内存在内核视野下：虚拟地址==物理地址
    // 注意这个映射关系只在内核的页目录表中存在存在，在进程的页目录表这个映射关系是没有的
    // 那么当我们的内核使用这个adrr开始的这页内存做内核栈，执行中断程序是，恰好就是用的内核自己的页目录表
    // 恰好就可以把addr这个物理地址当作虚拟地址使用
    uint32_t addr = addr_alloc_page(&paddr_alloc, 1);
    return addr;
}

static pde_t * curr_page_dir() {
    return (pde_t *)task_current()->tss.cr3;
}

// 释放一页内存，addr的可以是为内核分配的物理地址或者为精彩分配的物理地址相对应的虚拟地址
void memory_free_page(uint32_t addr) {
    // 由memory_alloc_page分配的物理内存，没有建立映射关系(这部分内存是被操作系统使用)
    if (addr < MEMORY_TASK_BASE) {
        addr_free_page(&paddr_alloc, addr, 1);
    } else {
        // 由memory_alloc_for_page_dir函数分配的物理内存，并且建立的虚拟内存和物理内存的映射关系(在页目录表和页表中)
        // 这部分内存是被进程所使用的
        // 进程空间，还要释放页表
        pte_t *pte = find_pte(curr_page_dir(), addr, 0);
        ASSERT((pte == (pte_t *)0) && pte->present);
        // 释放物理内存页
        addr_free_page(&paddr_alloc, pte_addr(pte), 1);
        // 释放页表
        pte->v = 0;
    }
}

void memory_destory_uvm(uint32_t page_dir) {
    // 找出用户空间页目录表项的起始索引
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    pde_t *pde = (pde_t *)page_dir + user_pde_start;
    // 遍历用户空间页目录项
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++) {
        if (!pde->present) {
            continue;
        }

        // 遍历页表
        pte_t *pte = (pte_t *)pde_addr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++) {
            if (!pte->present) {
                continue;
            }
            addr_free_page(&paddr_alloc, pte_addr(pte), 1);
        }
        addr_free_page(&paddr_alloc, (uint32_t)pde_addr(pde), 1);
    }
}

// 复制页表及其所有的内存空间
uint32_t memory_copy_uvm(uint32_t page_dir) {
    // 复制基础页表(内核的页表)
    uint32_t to_page_dir = memory_create_uvm();
    if (to_page_dir == 0) {
        goto copy_uvm_failed;
    }

    // 再复制用户空间的各项
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    pde_t *pde = (pde_t *)page_dir + user_pde_start;
    // 遍历用户空间页目录项
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++) {
        if (!pde->present) {
            continue;
        }

        // 遍历页表
        pte_t *pte = (pte_t *)pde_addr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++) {
            if (!pte->present) {
                continue;
            }
            // 分配物理内存
            uint32_t page = addr_alloc_page(&paddr_alloc, 1);
            if (page == 0) {
                goto copy_uvm_failed;
            }

            // 建立映射关系
            uint32_t vaddr = (i << 22) | (j << 12);
            int err = memory_create_map((pde_t *)to_page_dir, vaddr, page, 1, get_pte_perm(pte));
            if (err < 0) {
                goto copy_uvm_failed;
            }

            // 复制内存内容
            // 注意在内核视角下：物理内存页地址page == 虚拟地址page是相对应的
            // 这个拷贝函数看似简单实则很难理机
            // 首先我们就说使用这个memcpy函数主要是为了把当前父进程的某个页表中的第j项所指向的那个物理内存复制给子进程程序申请的一个物理页(物理地址page)
            // 我们观察参数page和vaddr，发现page是内存物理地址，vaddr则是父进程中的虚拟地址
            // 我们知道在程序中所用到的地址必须都是虚拟地址（在打开分页功能后）
            // 这样做是在内存总量不超过2GB时，因为父进程页目录表在创建页目录表时前512页（也就是2GB内存）是直接复制操作系统的
            // 而在这512页所映射的物理地址和虚拟地址是一一对应的（也就是前2GB内存是物理地址和虚拟地址是一一对应，可以互相混着用），
            // 因此对于物理地址page，在2GB内存以内，则可以把page当作虚拟的地址来使用，主要这里把虚拟地址page转换为物理地址page，使用的进程的页目录表中的前512项中的某一项来转换（其实本质这前512项就是内核页目录表的映射转换关系）
            // 而我们的vaddr变量它是进程中的真实的虚拟地址，它的地址>=0x8000 000 这个虚拟地址和物理地址的映射关系存放在当前进程(也就是父进程)页目录表的地513～1024项中的某个一项，
            // 而第513项到～1024项对于每个进程的映射关系都是不同的，不过无所谓，当我们访问这个vaddr时直接通过父进程的页目录表进行转换就能得到物理地址了
            // 
            kernel_memcpy((void *)page, (void *)vaddr, MEM_PAGE_SIZE);
        }
    }

    return to_page_dir;
copy_uvm_failed:
    if (to_page_dir) {
        memory_destory_uvm(to_page_dir);
    }
    return -1;
}

// 从当前页表curr->tss.cr3，在虚拟地址为frome，拷贝数量为size的情况放到新页目录表的to虚拟内存位置
int memory_copy_uvm_data(uint32_t to, uint32_t page_dir, uint32_t from, uint32_t size) {
    while (size > 0) {
        uint32_t to_paddr = memory_get_paddr(page_dir, to);
        if (to_paddr == 0) {
            return -1;
        }

        uint32_t offset_in_page = to_paddr & (MEM_PAGE_SIZE - 1);
        uint32_t curr_size = MEM_PAGE_SIZE - offset_in_page;
    
        if (curr_size > size) {
            curr_size = size;
        }
        // 主要这里必须每次一页一页的拷贝，因为to时物理地址，可能地址不是连续的！！！
        kernel_memcpy((void *)to_paddr, (void *)from, curr_size);
        
        size -= curr_size;
        to += curr_size;
        from += curr_size;
    }
}  


// 调整堆的内存分配，返回堆之前的指针
char * sys_sbrk(int incr) {
    task_t * task = task_current();
    char * pre_heap_end = (char * )task->heap_end;
    int pre_incr = incr;

    ASSERT(incr >= 0);

    // 如果地址为0，则返回有效的heap区域的顶端
    if (incr == 0) {
        // log_printf("sbrk(0): end = 0x%x", pre_heap_end);
        return pre_heap_end;
    } 
    
    uint32_t start = task->heap_end;
    uint32_t end = start + incr;

    // 起始偏移非0
    int start_offset = start % MEM_PAGE_SIZE;
    if (start_offset) {
        // 不超过1页，只调整
        if (start_offset + incr <= MEM_PAGE_SIZE) {
            task->heap_end = end;
            return pre_heap_end;
        } else {
            // 超过1页，先只调本页的
            uint32_t curr_size = MEM_PAGE_SIZE - start_offset;
            start += curr_size;
            incr -= curr_size;
        }
    }

    // 处理其余的，起始对齐的页边界的
    if (incr) {
        uint32_t curr_size = end - start;
        int err = memory_alloc_page_for(start, curr_size, PTE_P | PTE_U | PTE_W);
        if (err < 0) {
            log_printf("sbrk: alloc mem failed.");
            return (char *)-1;
        }
    }

    //log_printf("sbrk(%d): end = 0x%x", pre_incr, end);
    task->heap_end = end;
    return (char * )pre_heap_end;        
}